#ifndef TINYNDARRAY_ITERATOR_H
#define TINYNDARRAY_ITERATOR_H

namespace tinyndarray {

// ================================ 迭代器实现 ================================
template <bool Const>
class NdArray::IterBase {
public:
    using iterator_category = std::random_access_iterator_tag;
    using value_type = std::conditional_t<Const, const ElementProxy, ElementProxy>;
    using difference_type = std::ptrdiff_t;
    using pointer = value_type*;
    using reference = value_type&;

    IterBase(NdArray* array, size_t index) : array_(array), index_(index) {
        // 只在有效索引时更新代理
        if (index < array->size()) {
            update_proxy();
        }
    }

    // 解引用
    reference operator*() {
        if (index_ >= array_->size()) {
            throw std::out_of_range("Cannot dereference end iterator");
        }
        update_proxy();
        return proxy_;
    }

    // 成员访问
    pointer operator->() {
        if (index_ >= array_->size()) {
            throw std::out_of_range("Cannot dereference end iterator");
        }
        update_proxy();
        return &proxy_;
    }

    // 自增/自减
    IterBase& operator++() {
        ++index_;
        if (index_ < array_->size()) {
            update_proxy();
        }
        return *this;
    }

    IterBase operator++(int) {
        IterBase tmp = *this;
        ++index_;
        if (index_ < array_->size()) {
            update_proxy();
        }
        return tmp;
    }

    IterBase& operator--() {
        --index_;
        update_proxy();
        return *this;
    }

    IterBase operator--(int) {
        IterBase tmp = *this;
        --index_;
        update_proxy();
        return tmp;
    }

    // 算术运算
    IterBase operator+(difference_type n) const {
        return IterBase(array_, index_ + n);
    }

    IterBase operator-(difference_type n) const {
        return IterBase(array_, index_ - n);
    }

    IterBase& operator+=(difference_type n) {
        index_ += n;
        if (index_ < array_->size()) {
            update_proxy();
        }
        return *this;
    }

    IterBase& operator-=(difference_type n) {
        index_ -= n;
        update_proxy();
        return *this;
    }

    // 关系运算
    bool operator==(const IterBase& other) const {
        return array_ == other.array_ && index_ == other.index_;
    }

    bool operator!=(const IterBase& other) const {
        return !(*this == other);
    }

    bool operator<(const IterBase& other) const {
        return array_ == other.array_ && index_ < other.index_;
    }

    bool operator>(const IterBase& other) const {
        return array_ == other.array_ && index_ > other.index_;
    }

    bool operator<=(const IterBase& other) const {
        return array_ == other.array_ && index_ <= other.index_;
    }

    bool operator>=(const IterBase& other) const {
        return array_ == other.array_ && index_ >= other.index_;
    }

    // 转换操作符（用于const转换）
    operator IterBase<true>() const {
        return IterBase<true>(array_, index_);
    }

private:
    void update_proxy() {
        if (index_ >= array_->size()) {
            return;  // end() 位置不更新代理
        }
        size_t type_size = get_type_size(array_->dtype());
        void* ptr = reinterpret_cast<char*>(array_->m_sub->data.get()) + index_ * type_size;
        proxy_.reset(ptr, array_->dtype());
    }

    NdArray* array_;
    size_t index_;
    ElementProxy proxy_{nullptr, Dtype(FP32)};
};
}  // namespace tinyndarray

#endif